#!/usr/bin/env python
# coding: utf8
"""
Google appengine data model visualizer
by Maxim Vasiliev <qwiglydee@gmail.com>

inspired by modelviz by Antonio Cavedoni <antonio@cavedoni.org>

Put this script into appegine sdk directory and run from your app directory.
Listing modules to scan for models (by names as used to import)

$ modelviz.py [-h] [-d] [-g] [-i <model_names>] <module_name> ... <module_name> > <filename>.dot
$ dot <filename>.dot -Tpng -o <filename>.png

options:
    -h, --help
    show this help message and exit.

    -d, --disable_fields
    don't show the class member fields.

    -g, --group_modelsb
    draw an enclosing box around models from the same module.

    -i, --include_models=User,Person,Car
    only include selected models in graph.
"""

import getopt, sys, os
SDK_PATH = os.path.abspath(os.path.dirname(os.path.realpath(__file__)))
EXTRA_PATHS = [ '.', SDK_PATH, os.path.join(SDK_PATH, 'lib', 'webob') ]
sys.path = EXTRA_PATHS + sys.path

from google.appengine.ext.webapp.template import Template,Context
from google.appengine.ext import db

head_template = """
digraph name {
  fontname = "Helvetica"
  fontsize = 9
  pack = true

  node [
    fontname = "Helvetica"
    fontsize = 9
    shape = "plaintext"
  ]
  edge [
    fontname = "Helvetica"
    fontsize = 9
  ]
"""

body_template = """
{% if use_subgraph %}
subgraph cluster_{{ module_name }} {
  label=<
        <TABLE BORDER="0" CELLBORDER="0" CELLSPACING="0">
          <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER"><FONT FACE="Helvetica Bold" COLOR="Black" POINT-SIZE="12">{{ module_name }}</FONT></TD></TR>
        </TABLE>
        >
  color=dodgerblue4
  style="rounded"
{% endif %}

  {% for model in models %}
    {{ model.module_name }}_{{ model.name }} [label=<
    <TABLE BGCOLOR="darkslategray3" BORDER="0" CELLBORDER="0" CELLSPACING="0">
     <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="dodgerblue3">
       <FONT FACE="Helvetica Bold" COLOR="white">{{ model.name }}{% if model.abstracts %} (<FONT FACE="Helvetica">{{ model.abstracts|join:"," }}</FONT>){% endif %}
       </FONT>
     </TD></TR>

    {% if not disable_fields %}
        {% for field in model.fields %}
        <TR>
          <TD ALIGN="LEFT" BORDER="0">
            <FONT {% if not field.required %}COLOR="gray30" {% endif %}FACE="Helvetica {% if field.abstract %}Italic{% else %}Bold{% endif %}">{{ field.name }}</FONT>
          </TD>
          <TD ALIGN="LEFT">
            <FONT {% if not field.required %}COLOR="gray30" {% endif %}FACE="Helvetica {% if field.abstract %}Italic{% else %}Bold{% endif %}">{{ field.type }}</FONT>
          </TD>
        </TR>
        {% endfor %}
        {% if model.expando %}
        <TR><TD ALIGN="CENTER" BORDER="0" COLSPAN="2">
            <FONT COLOR="gray30" FACE="Helvetica Bold">...</FONT>
        </TD></TR>
        {% endif %}
    {% endif %}    
    </TABLE>
    >]
  {% endfor %}

{% if use_subgraph %}
}
{% endif %}
"""

rel_template = """
  {% for model in models %}
    {% for relation in model.relations %}
    {% if relation.needs_node %}
    {{ relation.target_module }}_{{ relation.target }} [label=<
        <TABLE BGCOLOR="darkslategray3" BORDER="0" CELLBORDER="0" CELLSPACING="0">
        <TR><TD COLSPAN="2" CELLPADDING="4" ALIGN="CENTER" BGCOLOR="dodgerblue3">
          <FONT FACE="Helvetica Bold" COLOR="white">{{ relation.target }}</FONT>
        </TD></TR>
        </TABLE>
        >]
    {% endif %}
    {{ model.module_name }}_{{ model.name }} -> {{ relation.target_module }}_{{ relation.target }}
    [label="{{ relation.name }}" headlabel="{{ relation.labels.1 }}" taillabel="{{ relation.labels.0 }}"] {{ relation.arrows }};
    {% endfor %}
  {% endfor %}
"""

import imp
def import_module(name):
    namepath = name.split('.')
    fullname = '_'.join(namepath)

    try:
        return sys.modules[fullname]
    except KeyError:
        pass

    if len(namepath) == 1:
        fp, pathname, description = imp.find_module(name,['.'])
    else:
        fp, pathname, description = imp.find_module(namepath[-1],[os.path.join(*namepath[:-1])])

    try:
        return imp.load_module(fullname, fp, pathname, description)
    finally:
        if fp:
            fp.close()

def generate_dot(module_names, **kwargs):
    disable_fields = kwargs.get('disable_fields', False)
    include_models = kwargs.get('include_models', [])
    use_subgraph = kwargs.get('group_models', False)

    dot = head_template

    graphs = []
    for name in module_names:
        module = import_module(name)

        graph = Context({
            'module_name': module.__name__,
            'disable_fields': disable_fields,
            'use_subgraph': use_subgraph,
            'models': []
        })

        for attr in dir(module):
            appmodel = getattr(module,attr)
            if not isinstance(appmodel,type) or not issubclass(appmodel,db.Model):
                continue
            
            def consider(model_name):
                return not include_models or model_name in include_models

            if not consider(appmodel.__name__):
                continue

            model =  {
                'module_name': appmodel.__module__,
                'name': appmodel.__name__,
                'abstracts': [],
                'fields': [],
                'relations': [],
                'expando': issubclass(appmodel,db.Expando)
            }

            def add_attribute(field):
                if isinstance(field,db.ListProperty):
                    typename = field.item_type.__name__ + '[]'
                else:
                    typename = field.data_type.__name__

                model['fields'].append({
                    'name': field.name,
                    'type': typename,
                    'required': field.required,
                    'abstract': False,
                })

            def add_relation(target,name='',arrows='',labels=''):
                rel = {
                    'target_module': target.__module__.replace('.','_'),
                    'target': target.__name__,
                    'name': name,
                    'labels': labels,
                    'arrows': arrows,
                    'needs_node': True
                }
                if rel not in model['relations'] and consider(rel['target']):
                    model['relations'].append(rel)

            def add_field(prop):
                add_attribute(prop)
                if hasattr(prop,'reference_class'):
                    if isinstance(prop,db.ListProperty):
                        labels = ['','*']
                    else:
                        labels = []                  
                    add_relation(prop.reference_class, prop.name, labels=labels)                

            if issubclass(appmodel,db.polymodel.PolyModel):
                bases = []
                for base in appmodel.__bases__:
                    if issubclass(base,db.polymodel.PolyModel) and base != db.polymodel.PolyModel:
                        bases.append(base)
                        add_relation(base, arrows='[arrowhead=onormal]')
                        model['abstracts'].append(base.__name__)
                for prop in appmodel._properties.itervalues():
                    if prop.name == 'class':
                        continue
                    inherited = False
                    for b in bases:
                        if hasattr(b,prop.name):
                            inherited = True
                            break
                    if not inherited:
                        add_field(prop)
            else:
                for prop in appmodel._properties.itervalues():
                    add_field(prop)

            graph['models'].append(model)
        graphs.append(graph)

    nodes = []
    for graph in graphs:
        nodes.extend([e['name'] for e in graph['models']])

    for graph in graphs:
        # don't draw duplication nodes because of relations
        for model in graph['models']:
            for relation in model['relations']:
                if relation['target'] in nodes:
                    relation['needs_node'] = False
        # render templates
        t = Template(body_template)
        dot += '\n' + t.render(graph)

    for graph in graphs:
        t = Template(rel_template)
        dot += '\n' + t.render(graph)

    dot += '\n}'
    return dot

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hadgi:",
                    ["help", "disable_fields", "group_models", "include_models="])
    except getopt.GetoptError, error:
        print __doc__
        sys.exit(error)
    
    kwargs = {}
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            print __doc__
            sys.exit()
        if opt in ("-d", "--disable_fields"):
            kwargs['disable_fields'] = True
        if opt in ("-g", "--group_models"):
            kwargs['group_models'] = True
        if opt in ("-i", "--include_models"):
            kwargs['include_models'] = arg.split(',')

    if not args:
        print __doc__
        sys.exit()

    print generate_dot(args, **kwargs)

if __name__ == "__main__":
    main()
